import math
import heapq
from collections import Counter

def knn_predict(
    x_test, 
    x_train, 
    y_train, 
    k=3, 
    distance_type='euclidean', 
    task='classification'
):
    """
    k近邻算法预测函数
    :param x_test: 单个测试样本（列表/元组，如[x1, x2, ...]）
    :param x_train: 训练集特征（列表嵌套列表，如[[x1,y1], [x2,y2], ...]）
    :param y_train: 训练集标签（列表，分类任务为离散值，回归任务为连续值）
    :param k: 近邻数量（正整数，默认3）
    :param distance_type: 距离度量类型（默认'euclidean'欧氏距离，可选'manhattan'曼哈顿距离）
    :param task: 任务类型（默认'classification'分类，可选'regression'回归）
    :return: 预测结果（分类为类别标签，回归为预测数值）
    """
    # 输入合法性检查
    if not x_train or not y_train:
        raise ValueError("训练集特征和标签不能为空")
    if len(x_train) != len(y_train):
        raise ValueError("训练集特征和标签数量必须一致")
    if k <= 0 or k > len(x_train):
        raise ValueError(f"k必须满足 0 < k ≤ 训练集数量（当前训练集数量：{len(x_train)}）")
    sample_dim = len(x_train[0])
    if len(x_test) != sample_dim:
        raise ValueError(f"测试样本维度（{len(x_test)}）与训练样本维度（{sample_dim}）不一致")
    if distance_type not in ['euclidean', 'manhattan']:
        raise ValueError("distance_type仅支持'euclidean'（欧氏距离）和'manhattan'（曼哈顿距离）")
    if task not in ['classification', 'regression']:
        raise ValueError("task仅支持'classification'（分类）和'regression'（回归）")
    
    # 1. 计算测试样本与所有训练样本的距离
    distances = []
    for i, x_train_sample in enumerate(x_train):
        if distance_type == 'euclidean':
            # 欧氏距离：d = sqrt(sum((x_test_i - x_train_i)^2))
            dist = math.sqrt(sum((test_dim - train_dim)**2 for test_dim, train_dim in zip(x_test, x_train_sample)))
        else:  # manhattan
            # 曼哈顿距离：d = sum(|x_test_i - x_train_i|)
            dist = sum(abs(test_dim - train_dim) for test_dim, train_dim in zip(x_test, x_train_sample))
        # 存储（距离，训练样本索引），方便后续取标签
        distances.append((dist, i))
    
    # 2. 找到距离最近的k个样本（使用小顶堆高效获取Top-K，时间复杂度O(n log k)）
    # heapq.nsmallest返回前k个最小元素，按距离升序排列
    k_nearest = heapq.nsmallest(k, distances, key=lambda x: x[0])
    
    # 3. 提取k个近邻的标签
    k_nearest_labels = [y_train[i] for (dist, i) in k_nearest]
    
    # 4. 根据任务类型计算预测结果
    if task == 'classification':
        # 分类任务：投票法（取出现次数最多的标签）
        counter = Counter(k_nearest_labels)
        # most_common(1)返回[(标签, 次数)]，取第一个元素的标签
        y_pred = counter.most_common(1)[0][0]
    else:  # regression
        # 回归任务：平均值法（取k个近邻标签的均值）
        y_pred = sum(k_nearest_labels) / len(k_nearest_labels)
    
    return y_pred

def knn_batch_predict(
    x_test_batch, 
    x_train, 
    y_train, 
    k=3, 
    distance_type='euclidean', 
    task='classification'
):
    """
    批量预测函数（对多个测试样本进行预测）
    :param x_test_batch: 测试集（列表嵌套列表，如[[x1,y1], [x2,y2], ...]）
    :return: 批量预测结果（列表）
    """
    return [
        knn_predict(x_test, x_train, y_train, k, distance_type, task)
        for x_test in x_test_batch
    ]